import torch
import torch.nn as nn
import sys
from pathlib import Path
CURRENT_DIR = Path(__file__).resolve().parent
if str(CURRENT_DIR) not in sys.path:
    sys.path.insert(0, str(CURRENT_DIR))
from 数据集加载器 import train_dl, val_dl, val_ds
from torchvision.models import alexnet, AlexNet_Weights

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ——加载预训练模型并冻结参数实现8分类———————————————————————————————————————————————————————
def load_alexnet(num_classes=8):
    model = alexnet(weights=AlexNet_Weights.IMAGENET1K_V1)
    # 冻结底层参数
    for p in model.parameters():
        p.requires_grad = False
    # 修改最后一层全连接 为8分类
    in_features = model.classifier[6].in_features
    out_features = num_classes
    model.classifier[6] = nn.Linear(in_features, out_features)
    return model.to(device)

model = load_alexnet()
# from torchsummary import summary
# summary(model, (3,224,224), 1)

# ———————————————————————————————————————————————————————————————————————————
lr = 0.001
Epochs = 10
loss_fn = nn.CrossEntropyLoss()
# 现在只训练最后一层
# 或者你也可以冻结特征提取器model.features、训练分类器model.classifier
optimizer = torch.optim.Adam(model.classifier[6].parameters(), lr=lr)
'''
我想要保存最优模型，采用在训练过程中监测 验证集的评估指标，去保存最优模型
'''
best_val_acc = 0.0
for epoch in range(Epochs):
    # 开启训练模式
    model.train()
    total_train_loss = 0
    for i, (imgs, labels) in enumerate(train_dl):
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        y_pre = model(imgs)
        train_loss = loss_fn(y_pre, labels)
        total_train_loss += train_loss.item()
        train_loss.backward()
        optimizer.step()
        if (i+1)%10==0:
            print(f"[{epoch+1}/{Epochs}] "
                  f"Batch [{i+1}/{len(train_dl)}]"
                  f"Loss {train_loss.item():.4f}")

    # 开启验证模式
    model.eval()
    total_val_acc = 0
    total_val_loss = 0
    with torch.no_grad():
        for imgs, labels in val_dl:
            imgs, labels = imgs.to(device), labels.to(device)
            y_pre = model(imgs)
            val_loss = loss_fn(y_pre, labels)
            total_val_loss += val_loss.item()
            y_pre_idx = y_pre.argmax(dim=-1)
            total_val_acc += (y_pre_idx == labels).sum().item()
    print(f"验证集 Loss:{total_val_loss/len(val_dl):.4f} Acc:{total_val_acc/len(val_ds)*100:.2f}%")

    # 进行判断，比假设更优需要更新迭代保存模型
    if total_val_acc > best_val_acc:
        best_val_acc = total_val_acc
        import os
        save_path = os.path.join(os.path.dirname(__file__), "best_alexnet.pth")
        torch.save(model.state_dict(), save_path)

print(f"最优验证集准确率为：{best_val_acc/len(val_ds)*100:.2f}%")

